/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// The main sfrhist code, except the actual calculation. \ingroup sfrhist

// $Id$

#include "config.h"
#include "hpm.h"
#include <unistd.h>
#include <time.h>
#include <vector>
#include <list>
#include <sstream>
#include <string>
#include "preferences.h"
#include "counter.h"
#include "CCfits/CCfits"
#include "misc.h"
#include "blitz-fits.h"
#include "model.h"
#include "sfrhist.h"
#include "sfrhist_snapshot.h"
#include "fits-utilities.h"
#include "boost/shared_ptr.hpp"
#include "boost/lambda/lambda.hpp"
#include "boost/lexical_cast.hpp"
#include "random.h"
#include <algorithm>
#include "sphgrid.h"
#include "grid-fits.h"
#include "mcrxversion.h"

#ifdef WITH_AREPO
#include "mpi.h"
#include "mcrx-arepo.h"
#endif

using namespace blitz;
using namespace std;
using namespace CCfits;
using namespace mcrx;
using namespace boost;
	   
/** Functor class that returns grid quantities at a specified
    location in an nbody_data_grid. This is supplied to the
    mappings_stellarmodel to pick the correct mappings model based on
    ISM pressure. */
class nbody_data_functor {
private:

  /// The grid that contains the pressure data.
  const nbody_data_grid& g_;

public:
  nbody_data_functor(const nbody_data_grid& g) : 
    g_(g) {};

  /** Returns the pressure/k (in cgs units!), gas density, and gas
      metallicity at the specified position. */
  pair<T_float,pair<T_float, T_float> > operator()(const vec3d& pos) const {
    const nbody_data_grid::T_cell* c = 
      const_cast<nbody_data_grid&>(g_).locate(pos).cell();

    if(c) {
      const T_float vol = c->volume();
      const T_float m_g = c->data()->m_g();
      const T_float density = 4.0463792e-8*m_g/vol; // density in mp/cm3
      const T_float temp = c->data()->gas_temp();
      const T_float teff = c->data()->gas_teff();
      const T_float mu = 1.21; // approx mol weight in mp (ign electrons)
      const T_float Pktemp = density*temp/mu; // P/k
      const T_float Pkteff = density*teff/mu; // P/k  
      const T_float Pk =Pkteff;//(use_teff_ ? Pkteff : Pktemp);
      assert(Pk >= 0);
      return make_pair(Pk, make_pair(m_g/vol, c->data()->m_metals()/m_g));
    }
    else {
      // particle is outside the grid, we have no pressure info
      // will this pick the lowest pressure we have?
      return make_pair(blitz::tiny(T_float()), make_pair(blitz::tiny(T_float()), 0.));
    }
  };
};


/** Reads TJ's merger set up file and puts the keywords into
    preferences object. */
bool mcrx::sfrhist::read_merger_file (const string & file_name) 
{
  ifstream file (file_name.c_str());
  if (!file) {
    cerr << "Error opening merger setup file: "
	 << file_name << endl;
    return false;
    //exit (1);
  }
  else
    cout << "Reading merger setup file: " << file_name << endl;
  
  // we read the file into a stream, looking for "GAL<i>"

  // keywords before is merger set up, after parameters for the
  // respective galaxy
  stringstream buffer_stream;
  string buffer;
  getline (file, buffer);
  while (file) {
    // put lines into buffer stream until we find a separator
    while (file && (buffer.substr (0,3) != "GAL") &&
	   (buffer.substr(0, 4) != "junk")) {
      buffer_stream << buffer << '\n';
      //cerr<< buffer <<endl;
      getline (file, buffer);
    }
    // now initialize a preferences object from the buffer stream
    merger.push_back(Preferences()); 
    merger.back().read(buffer_stream );
    buffer_stream.clear(); // clear EOF stream state after reading

    if (buffer.substr(0, 4) == "junk")
      // after this TJ puts junk we don't want
      break;
    
    getline (file, buffer);
  }
  cout << " The merger setup file contains " << merger.size()-1 << " objects"
       << endl;
  return true;
}

/* New order of initialization:
   1. load snapshot
   2. create sfrhist, assign z
   3. generate grid, save
   4. create model, which needs the grid
   5. assign SEDs, save particles
*/

/** Short function to open the snapshot file, load an enzo grid from
    it, and write output file. */
auto_ptr<nbody_data_grid> loadenzo(Preferences& p)
{
  cout << "Loading grid data from ENZO file" << endl;

  FITS file(p.getValue("snapshot_file", string()));
  ExtHDU& structure_hdu = open_HDU(file, "GRIDSTRUCTURE");
  ExtHDU& data_hdu = open_HDU(file, "GRIDDATA");
  vec3d translate_origin(0,0,0);
  if(p.defined("translate_origin")) {
    translate_origin = p.getValue("translate_origin", vec3d());
    cout << "  translating grid origin by " << translate_origin << endl;
  }
  auto_ptr<nbody_data_grid> g(new nbody_data_grid(structure_hdu,
						  translate_origin));
  cout << "\tgrid contains " << g->n_cells() << " cells" << endl;
  g->load_data(data_hdu);

  const string output_file = p.getValue("output_file", string());
  auto_ptr<FITS> output (new FITS("!"+output_file, Write));
  ExtHDU* makegrid_HDU = output->addTable("MAKEGRID", 0);
  makegrid_HDU->writeDate();
  makegrid_HDU->writeComment("Grid was imported from ENZO file.");
  // need to translate origin of grid data too

  cout << "Saving grid to file " << output_file << endl;
  g->save_structure(*output, "GRIDSTRUCTURE");
  // save length units in grid structure HDU
  g->save_data(*output, "GRIDDATA", "INTEGRATED_QUANTITIES");

  return g;
}


/// Setup global sfrhist stuff.
void global_setup(Preferences& p)
{

  // check if hpm should be enabled (default is disabled)
  if (p.defined("use_hpm") && p.getValue("use_hpm", bool ()))
    hpm::enable ();
  else
    hpm::disable ();

  //hpm::hpminit (0, "sfrhist2");

  // check if counters should be disabled (default is enabled)
  if (p.defined("use_counters") && !p.getValue("use_counters", bool ()))
    counter::enable = false;
  else
    counter::enable = true;

  // check if CCfits should be verbose (default is not)
  if (p.defined("CCfits_verbose"))
    CCfits::FITS::setVerboseMode (p.getValue("CCfits_verbose" , bool ()));

  // get seed and initialize the random number state
  const int seed = p.getValue("seed", int ());
  mcrx::seed (seed);

  p.setValue("pwd", string (getenv ("PWD")), "Working directory") ;
}


// decl of this in makegrid.cc
auto_ptr<nbody_data_grid> init_makegrid(Preferences& p, Snapshot& snap);


/** The sfrhist main function. This is very messy, there's no clear
    idea of what's being done by the sfrhist class, by the functions
    in makegrid and by functions here. It would be nice to clean this
    up...  \ingroup sfrhist */
int main(int argc, char** argv)
{
  cout << "sfrhist (Sunrise) version " << mcrx::version() << ". Copyright 2004-2011 Patrik Jonsson.\nThis program comes with ABSOLUTELY NO WARRANTY. Sunrise is free software, and\nyou are welcome to redistribute it under certain conditions. See the GNU\nGeneral Public License for details.\n\n";

  if(argc!=2) {
    cerr << "Usage:\n sfrhist <parameter file>\n";
    exit(1);
  }

  cout << "Welcome to The Stellar Model program, proudly SI compliant." 
       << endl << endl;

  assert(blitz::isThreadsafe());

  const string parameter_file = argv[1];
  Preferences prefs;
  if (!prefs.readfile (parameter_file)) {
    cerr << "Error opening parameter file: " << parameter_file << endl;
    exit (1);
  }
  else
    cout << "Reading parameter file: " << parameter_file << endl;
  // Add the name of the input file as a keyword
  prefs.setValue("configfile", parameter_file , "Configuration file") ;

  global_setup(prefs);

  // Read GADGET parameter file
  Preferences* gadget_prefs;
  if(prefs.defined("simparfile")) {
    const string simulation_parameter_file =
      word_expand(prefs.getValue ("simparfile", string ()))[0];
    const string simulation_directory =
      extract_directory (simulation_parameter_file);
    prefs.setValue("simdir", simulation_directory);
    gadget_prefs = new Preferences;
    if (!  gadget_prefs->readfile (simulation_parameter_file)) {
      cerr << "Error opening hydro simulation parameter file: \'"
	   << simulation_parameter_file << '\"'<<endl;
      exit (1);
    }
    else
      cout << "Reading hydro simulation parameter file: "
	   << simulation_parameter_file << endl;
  }
  else {
    // just link gadget prefs to main preferences object
    gadget_prefs = &prefs;
    cerr << "WARNING: No simulation parameter file specified\n  Simulation parameters will not be read.\n  Assuming relevant keywords are specified in config file." << endl;
  }
  
  // Get names of progenitor ic snapshots, if any are defined
  vector<string> progenitors;
  if(prefs.defined("ic_file1")) {
    const string ic_directory = word_expand(prefs.getValue("ic_snapshot_directory", string ()))[0];
    // extract names of progenitors until we throw
    try {
      int i=1;
      while(true)
	progenitors.push_back(ic_directory + 
			      prefs.getValue("ic_file"+
					     lexical_cast<string>(i++),string()));
    }
    catch (Preferences::unknown_keyword&) {}
  }
  cout << progenitors.size() << " IC snapshots defined." << endl;

  // load the snapshot
  const string snap_name = word_expand(prefs.getValue("snapshot_file", 
						  string ()))[0];

  sfrhist_snapshot snap(*gadget_prefs, prefs, snap_name, snap_name, 
			progenitors);

  mcrx::sfrhist s (prefs, *gadget_prefs, snap);

  s.assign_age_metallicity();

  // load stellar model
  auto_ptr<mcrx::stellarmodel> m;
  if(prefs.defined("mappings_sed_file")) {
    m.reset(new mcrx::mappings_sb99_model (prefs));
  }
  else
    m.reset(new mcrx::simple_stellarmodel (prefs));

  // create the grid, either by loading it or by projecting the
  // particles. (This also writes the grid quantities to the output
  // file.) IF we are using MAPPINGS, set the pressure_functor
  // appropriately).
  auto_ptr<nbody_data_grid> g;
  const string nbodycod = prefs.getValue("nbodycod",string());
  if (nbodycod == "enzo") {
    // if we are using enzo, we don't build a grid but just load it.
    g = loadenzo(prefs);
    m.get()->set_pressure_functor(nbody_data_functor(*g));
  }
  else if (nbodycod == "arepo") {
    // if we are using arepo, we actually don't build a grid at all
    // because the data is already in the arepo data structures
#ifdef WITH_AREPO
    g.reset();

    if(!prefs.defined("mappings_sed_file") || 
       !prefs.getValue("use_mappings_seds", true, true) ) {
      // if we are not using mappings, we don't need ism pressure so
      // we can skip loading snapshot with arepo now and just set a
      // dummy functor to return 0 for pressure and metallicity.
      m.get()->set_pressure_functor
	(boost::lambda::constant(make_pair(0.0,make_pair(0.,0.))));
      cout << "Not using MAPPINGS -- skipping loading Arepo now\n";
    }
    else {
      MPI_Init(&argc, &argv);
      arepo::load_snapshot(prefs.getValue("snapshot_file", string()), 
			   prefs.getValue("simparfile", string()), 
			   1, snap.units());
      m.get()->set_pressure_functor(arepo::arepo_pressure_functor());
    }
#else
    cerr << "Not compiled with Arepo support" << endl;
    exit(1);
#endif
  }
  else {
    g = init_makegrid(prefs, snap);
    m.get()->set_pressure_functor(nbody_data_functor(*g));
  }

  // Things that are now not done: radius assignment of all star particles

  // calculate creation mass of old stars
  snap.calculate_oldpop_massloss(*m);

  // now calculate the SEDs for star particles
  snap.calculate_stellar_SEDs(*m);

  auto_ptr<mcrx::bhmodel> b;
  if (prefs.defined("bhmodelfile")) {
    // Load black hole model
    b.reset(new mcrx::bhmodel(prefs));
    // Resample BH SEDs onto the wavelengths of the stellar SEDs
    b->resample( m->lambda () );
    
    // Now we can calculate the SEDs for the black hole particles.
    snap.calculate_BH_SEDs(*b);
  }

  // and finally bolometric luminosities
  snap.calculate_L_bol(*m);

  // and write output files
  // model writes STELLARMODEL, BHMODEL, and LAMBDA HDUs
  const string output_file_name =
    word_expand (strip_suffix (prefs.getValue("output_file", string ()))
		 + ".fits")[0];
  s.write_output_files (output_file_name);
  m->write_fits_parameters(output_file_name );
  if (b.get())
    b->write_fits_parameters(output_file_name );

  //hpm::hpmterminate (0);
}


mcrx::sfrhist::sfrhist (Preferences& pp, Preferences& gp,
			sfrhist_snapshot& ss) :
  p(pp), gadget(gp), snap(ss) {

  // Locate mergers/galaxy set up file, if applicable
  if(gadget.defined("InitCondFile")) {
    string merger_file =
      word_expand(gadget.getValue ("InitCondFile", string ()) +
		  ".parameters")[0]; 
    // decode this file into merger and galaxy set up info
    bool found_merger_file =read_merger_file (merger_file);
    if (!  found_merger_file) {
      // apparently, we couldn't open that file
      // see if we have defined an initial conditions directory
      if (p.defined("ic_directory")) {
	merger_file = word_expand(p.getValue("ic_directory", string ())+
				  strip_directory (merger_file))[0];
	found_merger_file = read_merger_file (merger_file);
      }
    }
    if (!found_merger_file) {
      cerr << "Warning: Initial conditions parameter file not found,\n  No initial conditions will be saved." <<  endl;
    }
  }
  else {
    cerr << "Warning: No initial conditions parameter file defined,\n  No initial conditions will be saved." <<  endl;
  }

  // set the is_merger bool based on the number of objects
  is_merger = merger.size()>0 ? (merger [0].size() > 0) : false;

  // check that the number of objects in the ic file is consistent
  // with the one in the snapshot  // is this REALLY strictly a failure?
  if (snap.nobject() != merger.size()-1) {
    cout << "Warning: The number of objects in the initial conditions file is different\n from the number of objects in the snapshot" << endl;
    //throw snap.nobject();
  }

  // Add some keywords that are going into HDU SFRHIST
  nobject = static_cast<int> ( merger.size() - 1);
  nobject = snap.nobject(); // THIS IS MORE RELIABLE
  p.setValue("NOBJECT", nobject, // also in the primary HDU
	     "Number of objects contained in this snapshot");
  

  // Calculate galaxy centers
  vector<vec3d> cm = snap.calculate_centers();
  
  mcrx::vec3d c(0.);
  for (int i = 0; i < nobject; ++i) {
    const string istr = lexical_cast<string> (i+1);
    p.setValue("galcenter" + istr, cm[i],
	       "["+snap.units().get("length")+"] Center of mass of galaxy "+ istr);
    c+=cm[i];
  }
  if(cm.size()>nobject)
    p.setValue("partcenter", cm[nobject],
	       "["+snap.units().get("length")+
	       "] Center of mass of free particles");

  c/= nobject;

  // if center_galaxies is set, translate mean of galaxy centers to origin
  if ( p.defined("center_galaxies") ||
       p.defined("translate_origin") ||
       p.defined("center_galaxy")) {
    if(p.defined("translate_origin"))
      c=p.getValue("translate_origin", vec3d());
    else if(p.defined("center_galaxy")) {
      const int gal=p.getValue("center_galaxy", int());
      if (gal<0) {
	if (cm.size()>nobject) {
	  c=cm.back();
	  cout << "Centering on free particles" << endl;
	}
	else
	  cerr << "Can't center on free particles -- none present." << endl;
      }
      if (gal>=0) {
	if(gal>=nobject)
	  cerr << "Can't center on galaxy " << gal << " -- not present" << endl;
	else {
	  c=cm[gal];
	  cout << "Centering galaxy " << gal << endl;
	}
      }
    }
    cout << "Translating origin to " << c << endl;
    if(p.getValue("nbodycod", string())!="arepo")
      // If we are using Arepo, we do NOT actually translate
      // anything. the cameras will be translated instead.
      snap.change_origin(c);
    p.setValue("translate_origin",c,
	       "["+snap.units().get("length")+"] Coordinates of new origin");
  }

  
  // This completes the initialization. Constructor ends here
}
  

/** Writes a history blurb to the specified HDU.  */
void mcrx::sfrhist::write_history (HDU & output)
{
  ostringstream history;
  history << "This file was created by Sunrise (sfrhist) version "
	  << mcrx::version() << ", on ";
  time_t now = time (0) ;
  char*now_string = ctime (& now);
  now_string [24] = 0;
  history << now_string << " by user ";
  const char* u=getenv ("USER");
  assert(u != 0);
  string user (u);
  const char* h=getenv ("HOST");
  if(h == 0)
    h = getenv ("HOSTNAME");
  if(h == 0)
    h = "unknown";
  assert(h != 0);
  string host (h);
  history << user << " on host "<< host << ".";
  output.writeHistory(history.str());
}
  

/** Write the output FITS files. Most of the writing is done by the
    snapshot object. */
void mcrx::sfrhist::write_output_files (const string& output_file_name)
{

  // okay, the main work is done.  Now all that remains is to output
  // all this stuff in the correct form to the correct files...
  cout << "Writing output file" << endl;

  // check for dump options (default is not dump)
  bool bolometric_dump = false;
  ofstream bolometric; 
  if (p.defined("bolometric_dump_file")) {
    bolometric_dump = true;
    bolometric.open(word_expand(p.getValue("bolometric_dump_file", 
					   string ()))[0].c_str ());
  }
  bool integrated_dump = false;
  ofstream integrated ; 
  if (p.defined("integrated_dump_file")) {
    integrated_dump = true;
    integrated.open(word_expand(p.getValue("integrated_dump_file", 
					   string ()))[0].c_str ());
  }
  
  cout << "  " << snap.name()<< endl;
  {
    // open local scope for output file
    FITS output_file (output_file_name, Write);
    
    // keywords that are going into the primary HDU
    output_file.pHDU().addKey("FILETYPE", string ("SNAPSHOT"),
			      "File contains a snapshot with computed SED");
    output_file.pHDU().addKey ("DATATYPE", string ("PARTICLE"),
			       "Data is a particle set");
    output_file.pHDU().addKey("NOBJECT", nobject,
			      "Number of objects contained in this snapshot");
    output_file.pHDU().addKey("SFRCODE", string("SFRHIST"),
			      "See HDU SFRHIST for SFR history calculation setup");
    output_file.pHDU().addKey("NBODYCOD", string("GADGET"),
			      "See HDU GADGET for nbody code setup");
    output_file.pHDU().addKey("L-INDEP", string("LAMBDA"),
			      "See HDU LAMBDA for SED wavelengths");
    output_file.pHDU().addKey("OBJECT",
			      is_merger? string ("MERGER"):string ("GALAXY"),
			      is_merger?  string ("See HDU MERGER for merger setup"):
			      string ("This snapshot contains an isolated object"));

    // output keywords to HDU "SFRHIST" (if we put it in the primary
    // HDU it's hard to copy)
    ExtHDU* sfrhist_HDU = output_file.addTable("SFRHIST", 0 );
    // write a bunch of info to the HISTORY field
    write_history (*sfrhist_HDU);
    sfrhist_HDU->writeDate();
    sfrhist_HDU->addKey("SNAPTIME", snap.time(),
			"["+snap.units().get("time")+
			"] Simulation time of this snapshot");
    sfrhist_HDU->addKey("SNAPFILE", snap.name(),
			"Snapshot file name");
    p.write_fits(*sfrhist_HDU);


    // GADGET HDU -- should really be more general depending on what
    // and body code we use
    ExtHDU* gadget_HDU = output_file.addTable("GADGET", 0);
    gadget.write_fits(*gadget_HDU );
    
    // MERGER HDU (if object is a merger)
    if (is_merger) {
      //cout << "    MERGER HDU" << endl;
      ExtHDU* merger_HDU = output_file.addTable("MERGER", 0);
      merger [0].write_fits(* merger_HDU );
    }

    // OBJECTi-PARAMETERS HDU (s)
    // in the case not all OBJECTs are present in the merger, we make sure we only write as many as are present.
    for (int i = 1; i < merger.size(); ++i) {
      ostringstream ost;
      ost << "OBJECT" << i << "-PARAMETERS";
      //cout << "    "<< ost.str() << " HDU" << endl;
      ExtHDU* object_HDU = output_file.addTable(ost.str(), 0);
      merger [i].write_fits(* object_HDU );
    }


  }

  //hpm::hpmtstart (80, "writing snapshot info");
  
  // and finally write the star particles
  snap.write_fits_table(output_file_name);
  
  //hpm::hpmtstop(80);
}
